package com.dl4jlpr;

import org.apache.commons.io.FilenameUtils;
import org.datavec.api.io.labels.PathLabelGenerator;
import org.datavec.api.writable.*;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

import java.io.*;
import java.net.URI;
import java.nio.charset.StandardCharsets;

public class PlateLableGenerator implements PathLabelGenerator {
    public static String CHARS = "京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新0123456789ABCDEFGHJKLMNPQRSTUVWXYZIO-";
    @Override
    public Writable getLabelForPath(String path) {
        char[] c=FilenameUtils.getBaseName(path).toCharArray();
        float[] l=new float[c.length];
        for(int i=0;i<c.length;i++)
        {
            l[i]=CHARS.indexOf(c[i]);
        }
        INDArray loadel=Nd4j.create(l,new int[]{1,c.length});
        return new NDArrayWritable(loadel);
    }

    @Override
    public Writable getLabelForPath(URI uri) {
        return this.getLabelForPath((new File(uri)).toString());
    }

    @Override
    public boolean inferLabelClasses() {
        return true;
    }
}
